import torch
import functools


CLS_LAYERS = 8
CLS_CHANNEL = 1024

VAN_ARGS = {
    'img_size': 512,
    'patch_size': 4,
    'in_chans': 3,
    'num_classes': 1000,
    'embed_dims': [64, 128, 320, 512],
    'num_heads': [1, 2, 5, 8],
    'mlp_ratios': [8, 8, 4, 4],
    'qkv_bias': True,
    'qk_scale': None,
    'drop_rate': .0,
    'attn_drop_rate': 0.0,
    'drop_path_rate': 0.1,
    'norm_layer': functools.partial(torch.nn.LayerNorm, eps=1e-06),
    'depths': [3, 3, 12, 3],
    'sr_ratios': [8, 4, 2, 1],
    'num_stages': 4,
    'linear': False,
}

UPER_ARGS = {
    'in_dims': [64, 128, 320, 512],
    'fc_dim': 4096,
    'fpn_dim': 256,
    'out_dim': 8,
    # 'pool_scales': [1, 2, 3, 6],
    'pool_scales': [1, 1, 2, 3],
}

FCN_ARGS = {
    'num_convs': 2,
    'kernel_size': 3,
    'concat_input': True,
    'dilation': 1,
    'in_channels': 320,
    'in_index': 2,
    'channels': 256,
    'dropout_ratio': 0.1,
    'num_classes': 8,
    'norm_cfg': {'type': 'BN', 'requires_grad': False},
    'align_corners': False,
    'loss_decode': {
        'type': 'CrossEntropyLoss',
        'use_sigmoid': False,
        'loss_weight': 0.4
    }
}
